package task

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"log"
	"m3u8d/util"
	"net/http"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/gookit/goutil/fsutil"
)

var (
	WorkerCount   = 100
	RetryCount    = 5
	TempTsFileDir = "./temp"
	HttpTimeOut   = time.Duration(30)
)

type DownloadItem struct {
	FileName string
	Url      string
}

// DownloadTsFile 下载ts文件
func DownloadTsFile(url string) (content []byte, err error) {
	client := http.Client{
		Timeout: HttpTimeOut * time.Second,
	}
	res, err := client.Get(url)
	if err != nil {
		return nil, err
	}
	defer res.Body.Close()
	content, err = ioutil.ReadAll(res.Body)
	if err != nil {
		return nil, err
	}
	return content, nil
}

// worker 下载文件的worker
func worker(wg *sync.WaitGroup, ch <-chan DownloadItem, done chan<- struct{}, info *util.M3u8Info) {
	for {
		item, ok := <-ch
		if !ok {
			break
		}
		filePath := fmt.Sprintf("%s/%s", TempTsFileDir, item.FileName)
		if !fsutil.FileExist(filePath) {
			content, err := DownloadTsFile(item.Url)
			if err != nil {
				for i := 0; i < RetryCount; i++ {
					time.Sleep(5 * time.Second)
					log.Println(item, "下载失败,重试次数：", i+1, err)
					content, err = DownloadTsFile(item.Url)
					if err == nil {
						break
					}
				}
			}
			if err != nil {
				log.Printf("%s重试%d次之后，下载失败，错误信息：%s\n", item, RetryCount, err)
				continue
			}
			if !bytes.Contains(content, []byte("FFmpeg")) && len(info.Key) > 0 {
				content, err = util.AesDecryptTsFile(content, info.Key, []byte(info.Iv))
				if err != nil {
					log.Println(item, "解密失败", err)
					continue
				}
			}
			if err = ioutil.WriteFile(filePath, content, os.ModePerm); err != nil {
				log.Println(item, "写入文件失败", err)
			}
		}
		done <- struct{}{}
	}
	wg.Done()
}

// DownloadItems 并发下载所有ts文件
func DownloadItems(url, outputFileName string) error {
	start := time.Now()
	wg := &sync.WaitGroup{}
	ch := make(chan DownloadItem)
	done := make(chan struct{})

	if !fsutil.DirExist(TempTsFileDir) {
		if err := os.MkdirAll(TempTsFileDir, os.ModePerm); err != nil {
			return err
		}
	}

	info, err := util.ParseM3u8File(url)
	if err != nil {
		return err
	}
	var items []DownloadItem
	for idx, url := range info.TsUrls {
		fileName := fmt.Sprintf("%d.ts", idx)
		items = append(items, DownloadItem{FileName: fileName, Url: url})
	}

	for i := 0; i < WorkerCount; i++ {
		wg.Add(1)
		go worker(wg, ch, done, info)
	}
	go ShowProgress(len(items), done)
	for _, item := range items {
		ch <- item
	}
	close(ch)
	// close(done)
	wg.Wait()
	fmt.Printf("下载完成，总共花费时间：%s\n", time.Since(start).String())
	outFileName := util.GetOrDefaultString(outputFileName, url[strings.LastIndex(url, "/"):strings.LastIndex(url, ".")]+".mp4")
	// 合并ts文件
	var fileNames []string
	for _, item := range items {
		fileNames = append(fileNames, item.FileName)
	}
	if err := util.ContactTsFilesDirectly(TempTsFileDir, fileNames, outFileName); err != nil {
		return err
	}
	// 删除 temp 文件夹
	if err := os.RemoveAll(TempTsFileDir); err != nil {
		return err
	}
	return nil
}

// ShowProgress 显示进度
func ShowProgress(all int, ch <-chan struct{}) {
	var count int
	for {
		if _, ok := <-ch; ok {
			count++
			if count == all {
				fmt.Println("\r下载进度：100.00%")
			} else {
				fmt.Printf("\r下载进度：%6.2f%%", float64(count*100)/float64(all))
			}
		} else {
			break
		}
	}
}
